#include <D3Dcompiler.h>
#include "ConstantStorageDX10.h"
#include <d3d10shader.h>
#pragma comment (lib, "D3DCompiler.lib")

CConstantStorageDX10::CConstantStorageDX10( )
{
    m_pBuffersData  = NULL;
    m_ppBuffers     = NULL;
    m_ppSRV         = NULL;
    m_pResData      = NULL;
    m_uNumResources = 0;
}

CConstantStorageDX10::~CConstantStorageDX10( )
{
    if (m_ppSRV)
    {
        /*for (mUInt uResInd = 0; uResInd < m_uNumResources; uResInd++)
        {
            SAFE_RELEASE(m_ppSRV[uResInd]);
        }*/
        delete [] m_ppSRV;
        delete [] m_pResData;
    }
    
    if (m_ppBuffers)
    {
        for (mUInt uBufInd = 0; uBufInd < ChangeFreqLast; uBufInd++)
        {
            if (m_pBuffersData[uBufInd].pData)
                delete [] m_pBuffersData[uBufInd].pData;
            SAFE_RELEASE(m_ppBuffers[uBufInd]);
        }        
    }
}

mVoid CConstantStorageDX10::Init( ID3D10Blob *a_pShaderCode )
{
    ID3D10ShaderReflection *pShaderReflection;

    /* Initialization of constants in each buffer: calculating its offset in memory and name */
    HRESULT hr;
    /* must be D3DReflect, but microsoft sucks balls =| */
    hr = D3D10ReflectShader(a_pShaderCode->GetBufferPointer(), a_pShaderCode->GetBufferSize(),
        /*IID_ID3D10ShaderReflection, (void**)*/&pShaderReflection);
    //hr == S_OK ?
    D3D10_SHADER_DESC ShaderDesc;
    D3D10_SHADER_INPUT_BIND_DESC ResDesc;
    pShaderReflection->GetDesc(&ShaderDesc);

    /* Check out for any resources - NOT CBUFFERS! */
    m_uNumResources = 0;
    UINT i;
    for (i = 0; i < ShaderDesc.BoundResources; i++)
    {
        pShaderReflection->GetResourceBindingDesc(i, &ResDesc);
        if (ResDesc.Type == D3D_SIT_TEXTURE)
            m_uNumResources++;
    }
    i = 0;
    if (m_uNumResources > 0)
    {
        m_pResData = new TResourceData[m_uNumResources];
        m_ppSRV    = new ID3D10ShaderResourceView*[m_uNumResources];
        ZeroMemory(m_ppSRV, m_uNumResources * sizeof(ID3D10ShaderResourceView*) );
        
        for (mUInt uResInd = 0; uResInd < ShaderDesc.BoundResources; uResInd++)
        {
            pShaderReflection->GetResourceBindingDesc(uResInd, &ResDesc);
            if (ResDesc.Type != D3D_SIT_TEXTURE)
                continue;            
            m_pResData[i++].sName = ResDesc.Name;
        }
    }

    /* Constant buffers */

    if (ShaderDesc.ConstantBuffers == 0)
    {
        SAFE_RELEASE(pShaderReflection);
        return;
    }

    m_pBuffersData = new TCBufferData [ChangeFreqLast];
    m_ppBuffers    = new ID3D10Buffer*[ChangeFreqLast];
    ZeroMemory(m_ppBuffers   , ChangeFreqLast * sizeof(ID3D10Buffer*));
    //ZeroMemory(m_pBuffersData, ChangeFreqLast * sizeof(TCBufferData ));


    /* 4 buffers - getting each by index and looking for name match */
    ID3D10ShaderReflectionConstantBuffer *pSRCB;
    D3D10_SHADER_BUFFER_DESC SRCBDesc;
    ID3D10ShaderReflectionVariable *pVarSRV;
    D3D10_SHADER_VARIABLE_DESC VarSVD;
    TConstantVariableData VarData;
    TCBufferData *pCurrentBufData;
    mUInt uVarOffset;
    mUInt uBufType;
    for (uBufType = 0; uBufType < ChangeFreqLast; uBufType++)
        m_pBuffersData[uBufType].pData = NULL;
    for (mUInt uBufInd = 0; uBufInd < ShaderDesc.ConstantBuffers; uBufInd++)
    {
        pSRCB = pShaderReflection->GetConstantBufferByIndex(uBufInd);
        pSRCB->GetDesc(&SRCBDesc);
        pCurrentBufData = m_pBuffersData;
        /* Looking for a match */
        for (uBufType = 0; uBufType < ChangeFreqLast; uBufType++)
        {
            if (g_sChangeFreqNames[uBufType] == SRCBDesc.Name)
                break;
            pCurrentBufData++;
        }
        /* No match - fuck off */
        if (uBufType == ChangeFreqLast)
        {
            continue;
        }
        pCurrentBufData->Size     = SRCBDesc.Size;
        pCurrentBufData->pData    = new mByte[SRCBDesc.Size];
        pCurrentBufData->bChanged = false;
        /* Since we have all info about buffer, 
        * let's calculating offset for each variable */
        uVarOffset = 0;
        for (mUInt uVarInd = 0; uVarInd < SRCBDesc.Variables; uVarInd++)
        {
            pVarSRV = pSRCB->GetVariableByIndex(uVarInd);
            pVarSRV->GetDesc(&VarSVD);
            VarData.Offset = uVarOffset;
            VarData.sName  = VarSVD.Name;
            pCurrentBufData->CVariables.push_back(VarData);
            uVarOffset += VarSVD.Size;
        }
    }    

    /* Initializing cbuffers in vram */
    D3D10_BUFFER_DESC cbDesc;
    cbDesc.ByteWidth        = 0;
    cbDesc.Usage            = D3D10_USAGE_DYNAMIC;
    cbDesc.BindFlags        = D3D10_BIND_CONSTANT_BUFFER;
    cbDesc.CPUAccessFlags   = D3D10_CPU_ACCESS_WRITE;
    cbDesc.MiscFlags        = 0;
    for (mUInt uBufInd = 0; uBufInd < ChangeFreqLast; uBufInd++)
    {
        if (m_pBuffersData[uBufInd].pData == NULL)
            continue;
        cbDesc.ByteWidth = m_pBuffersData[uBufInd].Size;
        g_D3DDevice.pD3D10Device->CreateBuffer(&cbDesc, NULL, &m_ppBuffers[uBufInd]);
    }
    /************************************************************/

    SAFE_RELEASE(pShaderReflection);
}

mBool CConstantStorageDX10::SendConstantValue( const std::string &a_sConstName, mVoid *a_pValue, size_t a_ValueSize, TChangeFreq a_ChangeFreq )
{
    if (m_ppBuffers == NULL)
        return false;
    if (m_ppBuffers[a_ChangeFreq] == NULL)
        return false;

    for (mUInt uVarInd = 0; uVarInd < m_pBuffersData[a_ChangeFreq].CVariables.size(); uVarInd++)
    {
        if (m_pBuffersData[a_ChangeFreq].CVariables[uVarInd].sName == a_sConstName)
        {
            mByte *pDst = (mByte*)m_pBuffersData[a_ChangeFreq].pData;
            mUInt uOffs = m_pBuffersData[a_ChangeFreq].CVariables[uVarInd].Offset;
            memcpy_s(pDst + uOffs, a_ValueSize, a_pValue, a_ValueSize);
            m_pBuffersData[a_ChangeFreq].bChanged = true;
            return true;
        }
    }
    //error! variable not found 
    return false;
}

mBool CConstantStorageDX10::AddResource( const std::string &a_sConstName, ID3D10ShaderResourceView *a_pSRV )
{
    if (m_pResData == NULL)
        return false;

    for (mUInt uResInd = 0; uResInd < m_uNumResources; uResInd++)
    {
        if (a_sConstName == m_pResData[uResInd].sName)
        {
            m_ppSRV[uResInd] = a_pSRV;
            return true;
        }
    }
    return false;
}

mVoid CConstantStorageDX10::UpdateConstantValues( )
{
    mVoid *pSrc;
    mVoid *pMappedBuf;
    if (m_ppBuffers == NULL)
        return;
    for (mUInt uBufInd = 0; uBufInd < ChangeFreqLast; uBufInd++)
    {
        if (m_ppBuffers[uBufInd] == NULL)
        {
            continue;
        }

        /* no changes - no update */
        if (!m_pBuffersData[uBufInd].bChanged)
            continue;

        m_ppBuffers[uBufInd]->Map( D3D10_MAP_WRITE_DISCARD, NULL, &pMappedBuf );
        size_t Size = m_pBuffersData[uBufInd].Size;
        pSrc = m_pBuffersData[uBufInd].pData;
        memcpy_s(pMappedBuf, Size, pSrc, Size);
        m_ppBuffers[uBufInd]->Unmap();
        m_pBuffersData[uBufInd].bChanged = false;
    }    
}